from threading import Thread
from queue import Queue, Empty
import time


class TimeoutError(Exception):
    pass


class WorkThread(Thread):
    def __init__(self, work_queue, timeout=None, **kwargs):
        super(WorkThread, self).__init__(kwargs=kwargs)
        self.work_queue = work_queue
        self.timeout = timeout
        self.setDaemon(True)
        self.END = False
        self.start()

    def run(self):
        while not self.END:
            func = args = result_queue = callback = None
            while not self.END:
                try:
                    func, args, result_queue, callback = self.work_queue.get(timeout=self.timeout)
                except Empty:
                    if self.END:
                        return
                else:
                    break
            try:
                result = func(*args)
                if callback:
                    callback(result)
                result_queue.put(result)
            except:
                raise


class ApplyResult:
    def __init__(self, result_queue):
        self.result_queue = result_queue

    def get(self, timeout=None):
        try:
            item = self.result_queue.get(timeout=timeout)
        except Empty:
            raise TimeoutError('timeout error')
        return item


class MapResult:
    def __init__(self, result_queue, apply_queue):
        self.result_queue = result_queue
        self.apply_queue = apply_queue

    def get(self, timeout=None):
        results = []
        endtime = 0
        if timeout is not None and timeout > 0:
            endtime = time.time() + timeout
        while not self.apply_queue.empty():
            remaining = endtime - time.time() if timeout is not None and timeout > 0 else None
            try:
                apply = self.apply_queue.get()
                result = apply.get(timeout=remaining)
            except:
                raise TimeoutError('timeout error')
            results.append(result)
        return results


class ThreadPool:
    def __init__(self, thread_num=5):
        self.work_queue = Queue()
        self.__threads = []
        self.thread_num = thread_num
        self.__create_thread_pool()
        self._closed = False

    def apply(self, func, args):
        assert not self._closed
        return self.apply_async(func, args).get()

    def apply_async(self, func, args, callback=None):
        assert not self._closed
        result_queue = Queue()
        self.__add_job(func, args, result_queue, callback)
        return ApplyResult(result_queue)

    def map(self, func, args):
        assert not self._closed
        return self.map_async(func, args).get()

    def map_async(self, func, args, callback=None):
        assert not self._closed
        result_queue = Queue()
        apply_queue = Queue()
        for arg in args:
            apply = self.apply_async(func, arg, callback)
            apply_queue.put(apply)
        return MapResult(result_queue, apply_queue)

    def imap(self, func, args):
        assert not self._closed
        return iter(map(func, args))

    def terminate(self):
        for thread in self.__threads:
            thread.END = True

    def close(self):
        pool.CLOSED = True

    def join(self):
        while self.__threads:
            work_thread = self.__threads.pop()
            if work_thread.isAlive():
                work_thread.join()

    def __create_thread_pool(self):
        for i in range(self.thread_num):
            work_thread = WorkThread(self.work_queue, timeout=0.1)
            self.__threads.append(work_thread)

    def __add_job(self, func, args, result_queue, callback=None):
        self.work_queue.put((func, args, result_queue, callback))


if __name__ == '__main__':
    def f(v):
        print(v)
        time.sleep(10)
        return v


    start = time.time()

    pool = ThreadPool(10)
    re = pool.map_async(f, [(i,) for i in range(50)])
    re.get(3)
    end = time.time()
    print(round((end - start), 2), 's')
    print(re)
